#!/usr/bin/env python
# -*- coding: utf-8 -*-
"""
UDP 转发模块 - 处理 UDP 连接转发
"""

import socket
import selectors
import time
import logging  # 添加这一行
from logger import setup_logger
from security import is_ip_allowed
from config import get_config_value

# 设置日志记录器
logger = setup_logger("udp_forwarder")

def udp_forward(local_port, remote_host, remote_port):
    """UDP 端口转发 - 优化版本"""
    try:
        # 创建 UDP 套接字
        server_socket = socket.socket(socket.AF_INET, socket.SOCK_DGRAM)
        # 设置接收和发送缓冲区大小
        server_socket.setsockopt(socket.SOL_SOCKET, socket.SO_RCVBUF, 65536)
        server_socket.setsockopt(socket.SOL_SOCKET, socket.SO_SNDBUF, 65536)
        server_socket.bind(('0.0.0.0', local_port))
        server_socket.setblocking(False)
        logger.info(f"UDP 转发已启动: 0.0.0.0:{local_port} -> {remote_host}:{remote_port}")
        
        # 创建 selector
        sel = selectors.DefaultSelector()
        sel.register(server_socket, selectors.EVENT_READ, None)
        
        # 客户端地址映射
        clients = {}
        remote_sockets = {}
        
        # 从配置中获取参数
        buffer_size = get_config_value("udp", "buffer_size", 8192)
        client_timeout = get_config_value("udp", "timeout", 120)  # 客户端超时时间（秒）
        max_clients = get_config_value("udp", "max_clients", 100)  # 最大客户端数量
        
        # 事件循环
        while True:
            events = sel.select(timeout=1)
            for key, mask in events:
                sock = key.fileobj
                
                if sock is server_socket:
                    # 处理来自客户端的数据
                    try:
                        data, client_addr = sock.recvfrom(buffer_size)
                        
                        # 检查客户端 IP 是否允许连接
                        if not is_ip_allowed(client_addr[0]):
                            logger.warning(f"拒绝来自 {client_addr[0]}:{client_addr[1]} 的 UDP 数据包 - IP 不在允许列表中")
                            continue
                        
                        # 检查是否已有此客户端的远程套接字
                        if client_addr not in clients:
                            # 检查是否达到最大客户端数量限制
                            if len(clients) >= max_clients:
                                logger.warning(f"拒绝来自 {client_addr[0]}:{client_addr[1]} 的 UDP 数据包 - 已达到最大客户端数量 {max_clients}")
                                continue
                                
                            # 为新客户端创建远程套接字
                            remote_socket = socket.socket(socket.AF_INET, socket.SOCK_DGRAM)
                            # 设置接收和发送缓冲区大小
                            remote_socket.setsockopt(socket.SOL_SOCKET, socket.SO_RCVBUF, 65536)
                            remote_socket.setsockopt(socket.SOL_SOCKET, socket.SO_SNDBUF, 65536)
                            remote_socket.setblocking(False)
                            
                            # 存储客户端信息
                            clients[client_addr] = {
                                'socket': remote_socket,
                                'addr': (remote_host, remote_port),
                                'last_active': time.time()
                            }
                            
                            # 注册远程套接字
                            sel.register(remote_socket, selectors.EVENT_READ, None)
                            remote_sockets[remote_socket] = client_addr
                            
                            logger.info(f"新的 UDP 客户端: {client_addr[0]}:{client_addr[1]} (当前客户端数: {len(clients)}/{max_clients})")
                        
                        # 更新最后活跃时间
                        clients[client_addr]['last_active'] = time.time()
                        
                        # 转发数据到远程服务器
                        clients[client_addr]['socket'].sendto(data, clients[client_addr]['addr'])
                        
                        # 记录数据传输（调试级别）
                        if logger.isEnabledFor(logging.DEBUG):
                            logger.debug(f"UDP 数据传输: 客户端->远程, 客户端: {client_addr[0]}:{client_addr[1]}, 大小: {len(data)} 字节")
                            
                    except BlockingIOError:
                        continue
                    except Exception as e:
                        logger.error(f"处理客户端 UDP 数据时出错: {e}")
                else:
                    # 处理来自远程服务器的响应
                    try:
                        data, _ = sock.recvfrom(buffer_size)
                        if sock in remote_sockets:
                            client_addr = remote_sockets[sock]
                            server_socket.sendto(data, client_addr)
                            clients[client_addr]['last_active'] = time.time()
                            
                            # 记录数据传输（调试级别）
                            if logger.isEnabledFor(logging.DEBUG):
                                logger.debug(f"UDP 数据传输: 远程->客户端, 客户端: {client_addr[0]}:{client_addr[1]}, 大小: {len(data)} 字节")
                    except BlockingIOError:
                        continue
                    except Exception as e:
                        logger.error(f"处理远程 UDP 响应时出错: {e}")
            
            # 清理超时客户端
            current_time = time.time()
            timeout_clients = []
            for client_addr, info in clients.items():
                if current_time - info['last_active'] > client_timeout:  # 使用配置的超时时间
                    timeout_clients.append(client_addr)
            
            for client_addr in timeout_clients:
                logger.info(f"UDP 客户端超时: {client_addr[0]}:{client_addr[1]} (超时时间: {client_timeout}秒)")
                remote_socket = clients[client_addr]['socket']
                sel.unregister(remote_socket)
                del remote_sockets[remote_socket]
                remote_socket.close()
                del clients[client_addr]
                
    except Exception as e:
        logger.error(f"UDP 转发错误: {e}", exc_info=True)
    finally:
        if 'server_socket' in locals():
            server_socket.close()
        if 'sel' in locals():
            sel.close()
        
        # 关闭所有客户端套接字
        for client in clients.values():
            client['socket'].close()